#! /usr/bin/ruby -w
#
# Agent script for Bliss project (SIGL p2007-p2008 @ INSIA).  Stops cleanly on
# SIGINT (Ctrl+C).
#
# @author Christophe Porteneuve [tdd@insia.org]
# @since  2007.05.31 10:35
# @latest 2007.06.04 17:32
#
# This script is update by Cyril Mougel for add :
#  * Define host with --host option
#  * If -r is an option there are notify random errors

require 'optparse'

module OS
  def self.windows?
    RUBY_PLATFORM =~ /win32/i
  end

  def self.osx?
    RUBY_PLATFORM =~ /darwin/i
  end
end

if OS.windows?
  STDOUT.puts "Works only with option -f and make fake"
end

# Interval between spontaneous status checks
CHECK_INTERVAL = 1 # In seconds

# Per-mountpoint available limit (KB).  Range must start with zero.
FREE_MP_KB_RANGE = 0..(50 * 1024)

# Available RAM limit (MB).  Range must start with zero.
FREE_RAM_MB_RANGE = 0..523

# URI to report incidents to.  Two params: kind and data.
# Expected to return a 201 Created with a Location header holding the update URI.
INCIDENT_URI = 'http://localhost:3000/incidents'

# List of mountpoints to check
if OS.windows?
  MOUNTPOINTS = %w(c:\\ e:\\ f:\\)
else
  MOUNTPOINTS = %w(/ /tmp /var /virtual)
end

# Port for local server (listening to explicit requests)
PORT = 1234

# There we go...

require 'net/http'
require 'socket'
require 'thread'
require 'uri'

ENV['LANG'] = 'C'

### Incident ############################

class Incident
  attr_reader :kind, :data

  def initialize(kind, data, uri_server)
    @kind = kind
    @data = data.to_s.split(' ').first
    @uri_server = uri_server
    @uri = 'not defined'
  end

  def report
    begin
      res = Net::HTTP.post_form @uri_server, { 'kind' => kind, 'data' => data }
      puts "res['Location'] = #{res['Location']}"
      @uri = res['Location']
      true
    rescue Exception => e
      STDERR.puts "#{Time.now} - #{e.class} while notifying incident (#{e.message})"
      false
    end
  end

  def resolve
    Net::HTTP.post_form URI.parse(@uri), { '_method' => 'put', 'resolved' => 'yes' }
  end

  def to_s
    @kind.to_s.upcase + ' ' + @data 
  end

  def ==(o)
    case self.kind
      when :lowram: return true
      when :lowmp:  return self.data == o.data
    end
  end
end # Incident

### StatusChecker #######################

class StatusSystem
  def self.check
    count = 0
    print Time.now.to_s + ' - Checking status...'
    free_ram_mb = get_ram_usage
  end

  def self.get_ram_usage
    if OS.osx?
      data = `vm_stat`
      return -1 if data.empty?
      data =~ /page size of (\d+) bytes.*Pages free:\s+(\d+)\./m
      return ($1.to_f / 1024**2 * $2.to_i).round
    else
      # UNIX/Linux
      data = `free -mt | cut -b30-40`
      return -1 if data.empty?
      data =~ /(\d+)\s+(\d+)/
      return $1.to_i + $2.to_i
    end
  end
  
  def self.get_mountpoint_usage(point)
    data = `df #{point} 2> /dev/null`
    return [-1, -1] if data.empty?
    data =~ /(\d+)\s+(\d+)%/
    return [$1.to_i, $2.to_i]
  end

end

class StatusChecker
  attr_reader :random
  attr_reader :uri
  attr_reader :force
  attr_reader :incidents

  def initialize (uri='http://localhost:3000/', random=false, force=false)
    @random = random
    @uri = uri
    @force = force
    @incidents = []
  end
  
  def run
    while true
      check
      sleep CHECK_INTERVAL
    end
  end

  def check
    count = 0

    if @random
      r = rand(10).round
      check_incident(false, :lowmp, 123) if r == 1
      check_incident(false, :lowram, 123) if r == 9
      check_incident(true, @incidents[0].kind, @incidents[0].data) if r == 5 && @incidents.size > 0
    else
      free_ram_mb = StatusSystem.get_ram_usage()
      count += check_incident(FREE_RAM_MB_RANGE.include?(free_ram_mb), :lowram, free_ram_mb, 0 == count)
      MOUNTPOINTS.each { |point|
        mp_usage = StatusSystem.get_mountpoint_usage(point)
        count += check_incident(FREE_MP_KB_RANGE.include?(mp_usage[0]), :lowmp, point + ' ' + mp_usage.join(' '))}
    end

    if count.zero?
      puts 'everything OK.'
    else
      puts '--> %d incident%s.' % [ count, count > 1 ? 's' : '']
    end
  end

  
private
  # if active is true, there are an error
  def check_incident(active, kind, data, first = false)
    incident = Incident.new kind, data, @uri
    unless active
      return 0 if @incidents.include?(incident) 
      text = kind.to_s.upcase + ' ' + data.to_s
      @incidents << incident if incident.report
    else
      puts "@incidents = #{@incidents}"
      inc_key = @incidents.index incident
      if inc_key
        puts "incident = #{incident}"
        puts "@incidents = #{@incidents}"
        @incidents[inc_key].resolve
        text = @incidents[inc_key].to_s + ' resolved'
        @incidents.delete incident
      else
        text = nil
      end
    end
    if text
      puts "\n" if first
      puts Time.now.to_s + ' - ' + text
    end
    active ? 1 : 0
  end
end # StatusChecker

### AgentServer #########################

class AgentServer
  def initialize
    @srv = TCPServer.new PORT
    
    trap("INT") do
      @srv.shutdown
      puts "Shutting down agent..."
      exit 0
    end
  end

  def run
    Thread.new do
      while (@clt = @srv.accept)
        begin
          while cmd = @clt.gets
	    process_command cmd.chomp
	  end
        rescue Exception => e
          STDERR.puts e unless IOError === e
          @clt.puts e unless @clt.closed?
        end
      end
    end
  end

private
  def process_command(cmd)
    case cmd
      when 'PING':            @clt.puts 'ALIVE'
      when 'RAM':             @clt.puts StatusSystem.get_ram_usage
      when /^MPUSE\s+(\S+)$/: @clt.puts StatusSystem.get_mountpoint_usage($1).join(' ')
      when /^\s*$/, 'QUIT':   @clt.close
      else                    @clt.puts "Invalid command: #{cmd}"
    end
  end
end # AgentServer

class CheckOption
  
  def self.parse(args)
    uri_server = nil
    random = false
    force = false
    options = {}
    begin
      opt = OptionParser.new do |opts|
        opts.banner = "Usage: #{opts.program_name} [options]"
        # View Help
        opts.on('-h', '--help', 'view this help') do |h|
          puts opts
          exit
        end

        opts.on('-u', '--host', 'host of Admin server', :REQUIRED) do |u|
          uri_server = URI::HTTP.build({:host => u,
                                         :path => '/incidents',
                                         :port => 3000})
        end

        opts.on('-r','--random', 'made random error') do |r|
          random = true
        end

        opts.on('-f', '--force', 'made fake result, works with windows') do 
          force = true
        end
      end
      opt.parse!(args)
    rescue
      puts opt
      exit
    end
    
    unless uri_server
      puts opt
      puts "--host required"
      exit
    end
    [random , uri_server, force]
  end
end

### MAIN PROGRAM ###

if __FILE__ == $0
  random, uri_server, force = CheckOption.parse ARGV
  AgentServer.new.run 
  sc = StatusChecker.new uri_server, random, force
  sc.run
end
